from torch.utils.data import DataLoader, Dataset
import torch
import numpy as np
import random

from .config import *
from .file_processor import FileProcessor

MAX_FEATURE_ROWS = 128  # 特征的最大行数


# 定义数据集
class SpeakerDataset(Dataset):
    def __init__(self, feature_dir):
        self.features = []
        self.labels = []

        def process_dataset(file_path, identity_id):
            feature = np.load(file_path)
            # 如果特征矩阵的列数不等于 FEATURE_LEN，则抛出异常
            if feature.shape[1] != featureLen:
                raise TypeError(f"Feature matrix shape error, except {featureLen}, get {feature.shape[1]}")

            # 分割特征矩阵为子样本点
            for i in range(0, feature.shape[0], MAX_FEATURE_ROWS):
                sub_feature = feature[i:i + MAX_FEATURE_ROWS]
                if sub_feature.shape[0] < MAX_FEATURE_ROWS:
                    # # 如果行数不够 MAX_FEATURE_ROWS，则进行零填充
                    # sub_feature = np.pad(sub_feature, ((0, MAX_FEATURE_ROWS - sub_feature.shape[0]), (0, 0)),
                    #                      mode='constant')
                    # 如果行数不够 MAX_FEATURE_ROWS，则舍弃
                    break
                # 将二维特征转换为三维
                sub_feature = sub_feature[np.newaxis, :, :]
                self.features.append(sub_feature)
                self.labels.append(int(identity_id[-4:]) - 1)

        fp = FileProcessor(feature_dir, process_dataset, output=False, max_amount=idAmount)
        fp.process_files()

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        feature = self.features[idx]
        label = self.labels[idx]
        return torch.tensor(feature, dtype=torch.float32), label


class SpeakerPairDataset(Dataset):
    def __init__(self, feature_dir):
        self.features = []
        self.labels = []
        self.speaker_to_features = {}  # 保存每个说话人的特征索引

        def process_dataset(file_path, identity_id):
            feature = np.load(file_path)
            # 检查特征矩阵的列数是否正确
            if feature.shape[1] != featureLen:
                raise TypeError(f"Feature matrix shape error, except {featureLen}, get {feature.shape[1]}")

            # 分割特征矩阵为子样本点
            for i in range(0, feature.shape[0], MAX_FEATURE_ROWS):
                sub_feature = feature[i:i + MAX_FEATURE_ROWS]
                if sub_feature.shape[0] == MAX_FEATURE_ROWS:  # 确保行数足够
                    sub_feature = sub_feature[np.newaxis, :, :]
                    index = len(self.features)
                    self.features.append(sub_feature)
                    self.labels.append(int(identity_id[-4:]) - 1)
                    if self.labels[-1] not in self.speaker_to_features:
                        self.speaker_to_features[self.labels[-1]] = []
                    self.speaker_to_features[self.labels[-1]].append(index)

        fp = FileProcessor(feature_dir, process_dataset, output=False, max_amount=idAmount)
        fp.process_files()

    def __len__(self):
        return len(self.features)

    def __getitem__(self, idx):
        feature1 = self.features[idx]
        label1 = self.labels[idx]

        # 选择第二个样本，使其与第一个样本同属一人或不同人
        if random.random() < 0.5:  # 同一个人的另一个样本
            if len(self.speaker_to_features[label1]) > 1:
                pair_idx = random.choice([i for i in self.speaker_to_features[label1] if i != idx])
            else:
                pair_idx = idx  # 没有其他样本可用，使用同一样本
        else:  # 不同人的样本
            different_label = random.choice([l for l in self.speaker_to_features if l != label1])
            pair_idx = random.choice(self.speaker_to_features[different_label])

        feature2 = self.features[pair_idx]
        label2 = self.labels[pair_idx]

        # 标签：1表示同一说话人，0表示不同说话人
        same_speaker = int(label1 == label2)

        return (torch.tensor(feature1, dtype=torch.float32),
                torch.tensor(feature2, dtype=torch.float32),
                torch.tensor(same_speaker, dtype=torch.float32))